import os
import pickle
import heapq
from collections import defaultdict, Counter
from typing import Tuple, Dict, List, Set

from pkg_resources import parse_requirements
from regex import T

from cs336_basics.bpe_tokenizer.pretokenizer import PreTokenizer


class Trainer:
    class heapData:
        """用于 heapq 进行排序的数据结构"""

        def __init__(self, pair: Tuple[bytes, bytes], freq: int):
            self.pair = pair
            self.freq = freq

        def __lt__(self, other):
            """定义优先级：频率高的优先，频率相同的，pair 字典序大的优先"""
            if self.freq != other.freq:
                return self.freq > other.freq
            return self.pair > other.pair

    def __init__(
        self,
        vocab_size: int,
        special_tokens: List[str] = [],
    ):
        self.vocab_size = vocab_size
        self.special_tokens = special_tokens or []
        self.vocab: Dict[int, bytes] = {}
        self.merges: List[Tuple[bytes, bytes]] = []
        self.pretokenizer = PreTokenizer(special_tokens=special_tokens)

        self.pair_freqs: Dict[Tuple[bytes, bytes], int] = {}  # 字节对出现的次数
        self.pair_in_pretokens: Dict[Tuple[bytes, bytes], Set[bytes]] = (
            {}
        )  # 字节对出现在哪些 pretoken 里
        self.pair_pos: Dict[bytes, Dict[Tuple[bytes, bytes], List[Tuple[int, int]]]] = (
            {}
        )  # 每个 pretoken 中，pair 出现的位置
        self.pretoken_indices: Dict[bytes, List[int]] = (
            {}
        )  # 每个 pretoken 的每个字节所属的 token 的 index

        self.heap: List[Trainer.heapData] = []  # 大根堆，按（频率，字典序）排序

    def _index_to_vocab(self, idx: int) -> bytes:
        """将索引转换为对应的词汇"""
        if idx in self.vocab:
            return self.vocab[idx]
        raise ValueError(f"Index {idx} not found in vocabulary.")

    def _pop_from_heap(self) -> Tuple[bytes, bytes]:
        """从大根堆中弹出堆顶的元素"""
        while self.heap:
            top = heapq.heappop(self.heap)
            if top.pair in self.pair_freqs and self.pair_freqs[top.pair] == top.freq:
                # 大根堆 lazy 删除，不一样说明这个状态的 pair 已被更新，需要删除
                return top.pair

        raise ValueError("Heap is empty, cannot pop from it.")

    def _push_to_heap(self, byte_pair: Tuple[bytes, bytes], freq: int):
        """将元素加入大根堆"""
        heapq.heappush(self.heap, Trainer.heapData(byte_pair, freq))

    def find_best_pair(self) -> Tuple[bytes, bytes]:
        """找到（频率，字典序）最大的 pair"""
        return self._pop_from_heap()

    def _update_pair(
        self,
        ptb: bytes,  # 当前的 pretoken
        ptb_freq: int,  # 当前 pretoken 的频次
        old_pair: Tuple[bytes, bytes],  # 旧的字节对
        old_pair_pos: Tuple[int, int],  # 旧的字节对在 pretoken 中的位置
        new_pair: Tuple[bytes, bytes],  # 新的字节对
        new_pair_pos: Tuple[int, int],  # 新的字节对在 pretoken 中的位置
    ):
        """更新字节对合并后的频次、所在pretoken集、pretoken内位置索引"""

        # 更新 old_pair 的频次
        if old_pair in self.pair_freqs:
            self.pair_freqs[old_pair] -= ptb_freq
            if self.pair_freqs[old_pair] <= 0:
                del self.pair_freqs[old_pair]
            else:
                self._push_to_heap(old_pair, self.pair_freqs[old_pair])

        # 删除 old_pair 在 pretoken 内的位置索引
        self.pair_pos[ptb][old_pair].remove(old_pair_pos)
        if len(self.pair_pos[ptb][old_pair]) == 0:
            del self.pair_pos[ptb][old_pair]

        # TODO：将 ptb 从 old_pair 的所在 pretoken 集中删除，提高后续查询性能，增强逻辑性与可读性
        # 当前不删除不会影响正确性，因为 merge 时会检查 pair 是否存在于 pretoken 中

        # 更新 new_pair 的频次
        self.pair_freqs[new_pair] = self.pair_freqs.get(new_pair, 0) + ptb_freq
        self._push_to_heap(new_pair, self.pair_freqs[new_pair])

        # 更新 new_pair 的所在 pretoken 集
        if new_pair not in self.pair_in_pretokens:
            self.pair_in_pretokens[new_pair] = set()
        self.pair_in_pretokens[new_pair].add(ptb)

        # 记录新 new_pair 的 pretoken 内位置索引
        if new_pair not in self.pair_pos[ptb]:
            self.pair_pos[ptb][new_pair] = []
        self.pair_pos[ptb][new_pair].append(new_pair_pos)

    def merge(
        self,
        ptb_freqs: Counter[bytes],
        pair: Tuple[bytes, bytes],
        new_idx: int,
    ):
        """合并 byte_pair"""

        # 找到 pair 在哪些 pretoken 中出现
        affected_pretokens = list(self.pair_in_pretokens.get(pair, set()))

        for ptb in affected_pretokens:
            len_pt = len(self.pretoken_indices[ptb])

            # 之所以加这个判断，是因为合并后并不会删除 pair_in_pretokens 中的 old_pair，所以有些 pair 可能在某些 pretoken 中并不存在
            if pair in self.pair_pos[ptb]:
                for l, r in self.pair_pos[ptb][pair]:
                    # 处理左边相邻的 token
                    if l > 0:
                        old_pair_l_pos = l - len(
                            self._index_to_vocab(self.pretoken_indices[ptb][l - 1])
                        )
                        old_pair_r_pos = (
                            l
                            + len(self._index_to_vocab(self.pretoken_indices[ptb][l]))
                            - 1
                        )
                        old_pair = (
                            self._index_to_vocab(self.pretoken_indices[ptb][l - 1]),
                            self._index_to_vocab(self.pretoken_indices[ptb][l]),
                        )
                        new_pair = (
                            self._index_to_vocab(self.pretoken_indices[ptb][l - 1]),
                            self._index_to_vocab(new_idx),
                        )

                        self._update_pair(
                            ptb,
                            ptb_freqs[ptb],
                            old_pair,
                            (old_pair_l_pos, old_pair_r_pos),
                            new_pair,
                            (old_pair_l_pos, r),
                        )

                    # 处理右边相邻的 token
                    if r < len_pt - 1:
                        old_pair_l_pos = (
                            r - len(self.vocab[self.pretoken_indices[ptb][r]]) + 1
                        )
                        old_pair_r_pos = r + len(
                            self.vocab[self.pretoken_indices[ptb][r + 1]]
                        )
                        old_pair = (
                            self._index_to_vocab(self.pretoken_indices[ptb][r]),
                            self._index_to_vocab(self.pretoken_indices[ptb][r + 1]),
                        )
                        new_pair = (
                            self._index_to_vocab(new_idx),
                            self._index_to_vocab(self.pretoken_indices[ptb][r + 1]),
                        )

                        self._update_pair(
                            ptb,
                            ptb_freqs[ptb],
                            old_pair,
                            (old_pair_l_pos, old_pair_r_pos),
                            new_pair,
                            (l, old_pair_r_pos),
                        )

                    # 更新 pair 的频次
                    if pair in self.pair_freqs:
                        self.pair_freqs[pair] -= ptb_freqs[ptb]
                        if self.pair_freqs[pair] <= 0:
                            del self.pair_freqs[pair]
                        else:
                            self._push_to_heap(pair, self.pair_freqs[pair])

                    # 更新 self.pretoken_indices
                    for i in range(l, r + 1):
                        self.pretoken_indices[ptb][i] = new_idx

                del self.pair_pos[ptb][pair]

    def train(self, input_path: str | os.PathLike, parallel_optimize: bool = False):
        """bpe 训练，使用 pretokenize，并优化 merge"""

        # 初始化 vocab 词汇表
        self.vocab = {x: bytes([x]) for x in range(256)}
        for idx, st in enumerate(self.special_tokens):
            self.vocab[256 + idx] = st.encode("utf-8")

        assert len(self.vocab) <= self.vocab_size

        # 初始化 merges 合并列表
        self.merges: list[tuple[bytes, bytes]] = []

        # 进行 pretokenization，并统计每个 pre_token 的频次
        pretokenizer = PreTokenizer(self.special_tokens)
        ptb_freqs = pretokenizer.count_pretoken_frequency(input_path, parallel_optimize)

        # 初始化 pretoken_indices
        self.pretoken_indices = {ptb: list(map(int, ptb)) for ptb in ptb_freqs}

        # 初始化各字节对出现的频次与位置索引
        self.pair_pos = {ptb: {} for ptb in ptb_freqs}
        for ptb in ptb_freqs:
            for pos_l in range(len(self.pretoken_indices[ptb]) - 1):
                pair = (
                    self._index_to_vocab(self.pretoken_indices[ptb][pos_l]),
                    self._index_to_vocab(self.pretoken_indices[ptb][pos_l + 1]),
                )

                self.pair_freqs[pair] = self.pair_freqs.get(pair, 0) + ptb_freqs[ptb]

                # 记录字节对出现在哪些 pretoken 中
                if pair not in self.pair_in_pretokens:
                    self.pair_in_pretokens[pair] = set()
                self.pair_in_pretokens[pair].add(ptb)

                # 记录字节对在 pretoken 中出现的位置
                if pair not in self.pair_pos[ptb]:
                    self.pair_pos[ptb][pair] = []
                self.pair_pos[ptb][pair].append((pos_l, pos_l + 1))

        # 初始化大根堆
        for pair, freq in self.pair_freqs.items():
            self._push_to_heap(pair, freq)

        # 开始合并
        base_idx = len(self.vocab)
        for i in range(self.vocab_size - len(self.vocab)):
            # 找到最佳的字节对
            pair = self.find_best_pair()
            bytes1, bytes2 = pair

            # 合并 token_pair
            new_idx = base_idx + i
            self.merges.append((bytes1, bytes2))
            self.vocab[new_idx] = bytes1 + bytes2
            self.merge(ptb_freqs, pair, new_idx)

    def get_vocab_and_merges(
        self,
    ) -> Tuple[Dict[int, bytes], List[Tuple[bytes, bytes]]]:
        """获取训练后的词汇表和合并规则"""
        return self.vocab, self.merges

    def save_to_file(
        self,
        vocab_filepath: str | os.PathLike,
        merges_filepath: str | os.PathLike,
    ):
        """保存词汇表和合并规则至文件"""
        with open(vocab_filepath, "wb") as f:
            pickle.dump(self.vocab, f)

        with open(merges_filepath, "wb") as f:
            pickle.dump(self.merges, f)
